use crate::{
    ast::{
        types::{BasicType as AstBasicType, MapType},
        Expression, Type,
    },
    error::{CompilerError, Result},
    llvm_backend::LLVMBackend,
};
use inkwell::{types::BasicType, values::BasicValueEnum, AddressSpace, IntPredicate};

impl<'ctx> LLVMBackend<'ctx> {
    /// Generate Map.contains(key) method call
    /// Returns true if key exists, false otherwise
    pub(in crate::llvm_backend) fn generate_map_contains_method(
        &mut self,
        var_name: &str,
        arguments: &[Expression],
    ) -> Result<BasicValueEnum<'ctx>> {
        if arguments.len() != 1 {
            return Err(CompilerError::internal(
                "Map.contains() requires exactly 1 argument (key)",
            ));
        }

        // Get the map variable pointer
        let map_ptr = self
            .variable_map
            .get(var_name)
            .ok_or_else(|| {
                CompilerError::internal(format!("Undefined map variable: {}", var_name))
            })?
            .into_pointer_value();

        // Get the map type
        let map_type = self.variable_type_map.get(var_name).ok_or_else(|| {
            CompilerError::internal(format!("Undefined map type for variable: {}", var_name))
        })?;

        let (key_type, _value_type) = if let Type::Map(MapType {
            key_type,
            value_type,
            ..
        }) = map_type
        {
            (*key_type.clone(), *value_type.clone())
        } else {
            return Err(CompilerError::internal(
                "contains() method requires a map variable",
            ));
        };

        // Generate the key expression
        let key_value = self.generate_expression(&arguments[0])?;

        // Get LLVM types
        let key_llvm_type = self.nrc_type_to_llvm_type(&Some(key_type.clone()))?;
        let value_llvm_type = self.nrc_type_to_llvm_type(&Some(_value_type.clone()))?;

        // Create map struct type
        let keys_ptr_type = key_llvm_type.ptr_type(AddressSpace::default());
        let values_ptr_type = value_llvm_type.ptr_type(AddressSpace::default());
        let map_struct_type = self.context.struct_type(
            &[
                keys_ptr_type.into(),
                values_ptr_type.into(),
                self.context.i64_type().into(),
                self.context.i64_type().into(),
            ],
            false,
        );

        // Load map
        let map_value = self
            .builder
            .build_load(map_struct_type, map_ptr, "map_value")?;
        let map_struct = map_value.into_struct_value();

        // Extract fields
        let keys_ptr = self
            .builder
            .build_extract_value(map_struct, 0, "keys_ptr")?
            .into_pointer_value();
        let len = self
            .builder
            .build_extract_value(map_struct, 2, "len")?
            .into_int_value();

        // Search for the key
        let key_index =
            self.find_key_index_for_contains(keys_ptr, len, key_value, &key_type, key_llvm_type)?;

        // Check if key was found (index >= 0)
        let zero = self.context.i64_type().const_int(0, false);
        let key_found =
            self.builder
                .build_int_compare(IntPredicate::SGE, key_index, zero, "key_found")?;

        // Return the boolean result
        Ok(key_found.into())
    }

    /// Find key index for contains method (returns -1 if not found)
    fn find_key_index_for_contains(
        &mut self,
        keys_ptr: inkwell::values::PointerValue<'ctx>,
        len: inkwell::values::IntValue<'ctx>,
        search_key: BasicValueEnum<'ctx>,
        key_type: &Type,
        key_llvm_type: inkwell::types::BasicTypeEnum<'ctx>,
    ) -> Result<inkwell::values::IntValue<'ctx>> {
        let current_func = self
            .current_function
            .ok_or_else(|| CompilerError::internal("No current function"))?;

        let loop_block = self
            .context
            .append_basic_block(current_func, "contains_find_loop");
        let body_block = self
            .context
            .append_basic_block(current_func, "contains_find_body");
        let found_block = self
            .context
            .append_basic_block(current_func, "contains_find_found");
        let not_found_block = self
            .context
            .append_basic_block(current_func, "contains_find_not_found");
        let merge_block = self
            .context
            .append_basic_block(current_func, "contains_find_merge");

        let entry_block = self.builder.get_insert_block().unwrap();

        self.builder.build_unconditional_branch(loop_block)?;

        // Loop block
        self.builder.position_at_end(loop_block);
        let i_phi = self.builder.build_phi(self.context.i64_type(), "i")?;
        let zero = self.context.i64_type().const_int(0, false);
        i_phi.add_incoming(&[(&zero, entry_block)]);
        let i = i_phi.as_basic_value().into_int_value();

        let loop_cond = self
            .builder
            .build_int_compare(IntPredicate::SLT, i, len, "loop_cond")?;
        self.builder
            .build_conditional_branch(loop_cond, body_block, not_found_block)?;

        // Body block
        self.builder.position_at_end(body_block);
        let key_ptr = unsafe {
            self.builder
                .build_gep(key_llvm_type, keys_ptr, &[i], "key_ptr")?
        };
        let current_key = self
            .builder
            .build_load(key_llvm_type, key_ptr, "current_key")?;

        let is_equal = self.compare_keys_for_contains(current_key, search_key, key_type)?;

        let i_next =
            self.builder
                .build_int_add(i, self.context.i64_type().const_int(1, false), "i_next")?;
        i_phi.add_incoming(&[(&i_next, body_block)]);

        self.builder
            .build_conditional_branch(is_equal, found_block, loop_block)?;

        // Found block
        self.builder.position_at_end(found_block);
        self.builder.build_unconditional_branch(merge_block)?;

        // Not found block
        self.builder.position_at_end(not_found_block);
        let minus_one = self.context.i64_type().const_int((-1i64) as u64, true);
        self.builder.build_unconditional_branch(merge_block)?;

        // Merge block
        self.builder.position_at_end(merge_block);
        let result_phi = self
            .builder
            .build_phi(self.context.i64_type(), "find_result")?;
        result_phi.add_incoming(&[(&i, found_block), (&minus_one, not_found_block)]);

        Ok(result_phi.as_basic_value().into_int_value())
    }

    /// Compare keys for contains method
    fn compare_keys_for_contains(
        &mut self,
        key1: BasicValueEnum<'ctx>,
        key2: BasicValueEnum<'ctx>,
        key_type: &Type,
    ) -> Result<inkwell::values::IntValue<'ctx>> {
        match key_type {
            Type::Basic(basic_type) => {
                match basic_type {
                    AstBasicType::Int8
                    | AstBasicType::Int16
                    | AstBasicType::Int32
                    | AstBasicType::Int
                    | AstBasicType::Int64
                    | AstBasicType::Uint8
                    | AstBasicType::Uint16
                    | AstBasicType::Uint32
                    | AstBasicType::Uint
                    | AstBasicType::Uint64
                    | AstBasicType::Bool => {
                        let int1 = key1.into_int_value();
                        let int2 = key2.into_int_value();

                        // Handle type mismatch by casting to the wider type
                        let (cmp1, cmp2) = if int1.get_type() != int2.get_type() {
                            let width1 = int1.get_type().get_bit_width();
                            let width2 = int2.get_type().get_bit_width();

                            if width1 > width2 {
                                let extended = self.builder.build_int_z_extend_or_bit_cast(
                                    int2,
                                    int1.get_type(),
                                    "extend",
                                )?;
                                (int1, extended)
                            } else if width2 > width1 {
                                let extended = self.builder.build_int_z_extend_or_bit_cast(
                                    int1,
                                    int2.get_type(),
                                    "extend",
                                )?;
                                (extended, int2)
                            } else {
                                (int1, int2)
                            }
                        } else {
                            (int1, int2)
                        };

                        Ok(self.builder.build_int_compare(
                            IntPredicate::EQ,
                            cmp1,
                            cmp2,
                            "key_eq",
                        )?)
                    }
                    AstBasicType::String => {
                        let strcmp_func = self
                            .function_map
                            .get("strcmp")
                            .ok_or_else(|| CompilerError::internal("strcmp function not found"))?;

                        let strcmp_result = self.builder.build_call(
                            *strcmp_func,
                            &[
                                key1.into_pointer_value().into(),
                                key2.into_pointer_value().into(),
                            ],
                            "strcmp_result",
                        )?;

                        let result_int = strcmp_result
                            .try_as_basic_value()
                            .left()
                            .unwrap()
                            .into_int_value();
                        let zero = self.context.i32_type().const_int(0, false);

                        Ok(self.builder.build_int_compare(
                            IntPredicate::EQ,
                            result_int,
                            zero,
                            "strings_equal",
                        )?)
                    }
                    _ => Err(CompilerError::internal(
                        "Unsupported key type for comparison",
                    )),
                }
            }
            _ => Err(CompilerError::internal(
                "Only basic types supported as map keys",
            )),
        }
    }
}
